/**
 * @file Predict.cpp
 * @author enemy1205 (enemy1205@qq.com)
 * @brief 预测类实现
 * @date 2021-10-10
 */
#include "Predict.h"
#include <Eigen/Core>
#include<Eigen/Geometry>
#include <opencv2/core/eigen.hpp>
#include <opencv2/calib3d/calib3d.hpp>
//------------------------基类函数---------------------

/**
 * @brief 利用距离卡尔曼更新pnp数据
 * @param tracker 待预测追踪器对象
 */
void Predict::updatePnP(tracker_ptr &tracker) {
    if (tracker->getGoalFeature() == nullptr) return;
    tracker->setPnp() = calculatePnPData(cameraParam.cameraMatrix, cameraParam.distCoeff, tracker);
    if (tracker->getDisKalmanInfo()) {
        // 设置状态转移矩阵
        tracker->setDistanceKF().setA(Matx<float, 2, 2>{1, 0,
                                      0, 1});
        float delta_distance = tracker->getPnp().distance - last_distance;
        // 预测
        tracker->setDistanceKF().predict();
        // 更新
        auto correctVec = tracker->setDistanceKF().correct({tracker->getPnp().distance,
                                          delta_distance});
        tracker->setPnp().distance = correctVec(0);
    } else {
        cv::Matx21f initVec = cv::Matx21f(tracker->getPnp().distance, 0);
        tracker->setDistanceKF().init(initVec, 1e-2);
        tracker->setDisKalmanInfo() = true;
    }
    last_distance = tracker->getPnp().distance;
}

//--------------------------装甲板预测部分-------------------------

/**
 * @brief 更新记录数据序列
 * @param gyroData 电控发来的陀螺仪数据
 * @param appear_time 识别到装甲板的时间
 */
void ArmorPredict::updateSequence(GyroData &gyroData, clock_t &appear_time) {
    //-----------------------记录数据----------------------------
    this->_gyroData.emplace_front(make_unique<GyroData>(gyroData));
    this->time_sequence.emplace_front(appear_time);
    //只保留四帧电控数据
    if (this->_gyroData.size() > 3) this->_gyroData.pop_back();
    // 保持时间戳最大长度
    if (this->time_sequence.size() >= (size_t) A_param.Seq_maxlength) {
        // 从后 pop
        this->time_sequence.pop_back();
    }
}

/**
 * @brief 装甲板预测初始化
 */
ArmorPredict::ArmorPredict() {
    //-----------------------参数初始化----------------------------
    this->last_distance = 0.f;
    this->Speed = cv::Point2f(0, 0);
    this->target_speed = cv::Point2f(0, 0);
    this->gyro_speed = cv::Point2f(0, 0);
}

/**
 * @brief 计算pnp数据
 * @param cameraMatrix 相机内参
 * @param distcoeff 畸变矩阵
 * @param tracker 待预测追踪器对象
 * @return pnp结果
 */
ResultPnP ArmorPredict::calculatePnPData(const Mat &cameraMatrix, const Mat &distcoeff, tracker_ptr &tracker) {
    // 像素坐标系下的坐标点
    vector<Point2f> img_points;
    img_points.emplace_back(tracker->getGoalFeature()->getVertexes()[0]);//左下
    img_points.emplace_back(tracker->getGoalFeature()->getVertexes()[1]);//左上
    img_points.emplace_back(tracker->getGoalFeature()->getVertexes()[2]);//右上
    img_points.emplace_back(tracker->getGoalFeature()->getVertexes()[3]);//右下
    Mat rvec;              // 旋转向量
    Mat tvec;              // 平移向量
    ResultPnP pnp_message; // 存储 PNP 信息
    //根据大小装甲板分别算pnp
    if (tracker->getGoalFeature()->wh_ratio() < A_param.BIG_armor_min_ratio) {
        solvePnP(A_param.Small_Armor, img_points, cameraMatrix, distcoeff, rvec, tvec, false, cv::SOLVEPNP_IPPE);
    } else {
        solvePnP(A_param.Big_Armor, img_points, cameraMatrix, distcoeff, rvec, tvec, false, cv::SOLVEPNP_IPPE);
    }
    Mat rmat; // 旋转矩阵
    // 罗德里格斯变换: 旋转向量 -> 旋转矩阵
    Rodrigues(rvec, rmat);
    rmat.convertTo(rmat, CV_64FC1);
    tvec.convertTo(tvec, CV_64FC1);
    Eigen::Matrix3d rotated_matrix;  // 旋转矩阵
    Eigen::Vector3d transfer_vector; // 平移向量
    // 类型转换
    cv2eigen(rmat, rotated_matrix);
    cv2eigen(tvec, transfer_vector);
    // 获取欧拉角
    Eigen::Vector3d euler_angles = rotated_matrix.eulerAngles(0, 1, 2);
    pnp_message.pitch = Rad2Deg(euler_angles[0]);
    pnp_message.yaw = Rad2Deg(euler_angles[1]);
    pnp_message.roll = Rad2Deg(euler_angles[2]);
    pnp_message.distance =
            cosf(3 * Deg2Rad(tracker->getGoalFeature()->getAngle())) *
            sqrt(transfer_vector.transpose() * transfer_vector);
//    if (this->armor_type == ArmorType::SMALL) {
//        pnp_message.distance = cosf(A_param.SMALL_COEFF_W * Deg2Rad(this->angle)) *
//                               sqrt(transfer_vector.transpose() * transfer_vector);
//    } else {
//        pnp_message.distance = cosf(A_param.BIG_COEFF_W * Deg2Rad(this->angle)) *
//                               sqrt(transfer_vector.transpose() * transfer_vector);
//    }
    tvec.convertTo(tvec, CV_32FC1);
    pnp_message.tran_vec = tvec;
    pnp_message.rotat_vec = rvec;
    return pnp_message;
}

/**
 * @brief 根据陀螺仪数据获取陀螺仪速度
 */
void ArmorPredict::getGyroSpeed() {
    //原本想用近三帧加权
/*    if (gyroData.size() > 2) {
        gyro_speed.x = gyroData[0]->yawSpeed / 920;
        gyro_speed.y = gyroData[0]->pitchSpeed / 920;
    } else */
    gyro_speed = cv::Point2f(_gyroData[0]->yawSpeed, _gyroData[0]->pitchSpeed) / 920;
}

/**
 * @brief 为装甲板更新卡尔曼信息
 * @param times 时间序列
 */
void ArmorPredict::updateKalman(tracker_ptr &tracker, float flying_time) {
    if (tracker->getKalmanInfo()) {
        float predict_yaw_time = flying_time;
        float predict_pitch_time = flying_time;
        // 状态转移矩阵
        tracker->setArmorKF().setA(cv::Matx44f{
                1, 0, predict_yaw_time, 0,
                0, 1, 0, predict_pitch_time,
                0, 0, 1, 0,
                0, 0, 0, 1});
        // 不出现闪烁
        if (!tracker->isShine()) {
            // 帧差时间 -- init
            float total_time;
            // 对临近数帧帧差值取平均
            if (time_sequence.size() >= 4 && tracker->getFeatures().size() >= 4) {
                total_time = 0.1f + static_cast<float >(time_sequence[0] + time_sequence[1] - time_sequence[2] -
                                                        time_sequence[3]) / (4 * CLOCKS_PER_SEC);
                // 相对角速度
                this->target_speed = (tracker->getGoalFeature()->getRelativeAngles() +
                                      tracker->getFeatures()[1]->getRelativeAngles() -
                                      tracker->getFeatures()[2]->getRelativeAngles() -
                                      tracker->getFeatures()[3]->getRelativeAngles()) /
                                     (2.f * total_time);
//                this->getGyroSpeed();
            } else if (time_sequence.size() >= 2 && tracker->getFeatures().size() >= 2) {
                this->target_speed = (tracker->getGoalFeature()->getRelativeAngles() -
                                      tracker->getFeatures()[1]->getRelativeAngles()) /
                                     (static_cast<float >(time_sequence[0] - time_sequence[1]) / (2 * CLOCKS_PER_SEC));
//                this->getGyroSpeed();
            }
            // 绝对速度
            this->Speed = this->target_speed + this->gyro_speed;
//            cout << "target_speed = " << this->target_speed << endl;
//            cout << "gyro_speed = " << this->gyro_speed << endl;
//            cout << "speed = " << this->Speed << endl;
            // 更正测量值
            tracker->setArmorKF().correct({tracker->getGoalFeature()->getRelativeAngles().x, // 位置_x
                              tracker->getGoalFeature()->getRelativeAngles().y, // 位置_y
                              this->Speed.x,                               // 速度_x
                              this->Speed.y});                             // 速度_y
        }
        // 预测
        auto predict_mat = tracker->setArmorKF().predict();
        // 更新卡尔曼预测值
        this->updatePredict(predict_mat, tracker);
    } else {
        // 初始化为装甲板实际中心
        cv::Matx41f initMat(tracker->getGoalFeature()->getRelativeAngles().x,
                            tracker->getGoalFeature()->getRelativeAngles().y,
                            0, 0);
        // 初始化卡尔曼滤波器
        tracker->setArmorKF().init(initMat, 1e-5);
        // 更新卡尔曼预测值
        this->updatePredict(initMat, tracker);
        // 更新标志位
        tracker->setKalmanInfo() = true;
    }
}

/**
 * @brief 根据卡尔曼信息更新预测值
 * @param predict_mat 预测信息矩阵
 */
void ArmorPredict::updatePredict(const cv::Matx41f &predict_mat, tracker_ptr &tracker) {
    // 更新预测角度
    tracker->setPredictedAngle().x = predict_mat(0);
    tracker->setPredictedAngle().y = predict_mat(1);
    // 更新预测中心点
    tracker->setPredictedCenter() = (1 - A_param.predictConfidence) * tracker->getGoalFeature()->getCenter() +
                                    A_param.predictConfidence * calculateRelativeCenter(cameraParam.cameraMatrix,
                                                                                        cameraParam.distCoeff,
                                                                                        tracker->getPredictedAngle());
}

void ArmorPredict::predict(tracker_ptr &tracker, Compensate_ptr &compensate) {
    this->updatePnP(tracker);//为正在跟踪的每个tracker更新pnp信息
    compensate->compensate(tracker, _gyroData.front());
    this->updateKalman(tracker, compensate->getTime());//为正在跟踪的每个tracker更新kalman信息

}


//-------------------------------能量机关部分--------------------------------

/**
 * @brief 更新记录数据序列
 * @param gyroData 电控发来的陀螺仪数据(神符并不需要，但要在基类为装甲板识别提供接口)
 * @param appear_time 识别到装甲板的时间
 */
void RunePredict::updateSequence(GyroData &gyroData, clock_t &appear_time) {
    //-----------------------记录数据----------------------------
    this->time_sequence.emplace_front(appear_time);
    // 保持时间戳最大长度
    if (this->time_sequence.size() >= (size_t) A_param.Seq_maxlength) {
        // 从后 pop
        this->time_sequence.pop_back();
    }
}

/**
 * @brief 神符预测构造函数初始化
 */
RunePredict::RunePredict() {
    //-----------------------参数初始化----------------------------
    this->last_distance = 0.f;
    this->image_predict_angle = 0.f;
    this->radius = 0.f;
    this->speed = 0.f;
    this->round = 0;
}

/**
 * @brief 计算pnp数据
 * @param cameraMatrix 相机内参
 * @param distcoeff 畸变矩阵
 * @param tracker 待预测追踪器对象
 * @return pnp结果
 */
ResultPnP RunePredict::calculatePnPData(const Mat &cameraMatrix, const Mat &distcoeff, tracker_ptr &tracker) {
    // 像素坐标系下的坐标点,依次：左下，左上，右上，右下，中点
    vector<Point2f> img_points(5, Point2f(0, 0));
    for (int i = 0; i < 4; ++i) {
        img_points[i] = tracker->getGoalFeature()->getSon1_Features()->getVertexes()[i];
    }
    img_points[4] = tracker->getGoalFeature()->getSon1_Features()->getCenter();
    Mat rvec;              // 旋转向量
    Mat tvec;              // 平移向量
    ResultPnP pnp_message; // 存储 PNP 信息
    // 姿态解算
    solvePnP(R_param.wRUNE_CONTOUR, img_points, cameraMatrix, distcoeff, rvec, tvec, false, cv::SOLVEPNP_IPPE);
    Mat rmat; // 旋转矩阵
    // 罗德里格斯变换: 旋转向量 -> 旋转矩阵
    Rodrigues(rvec, rmat);
    rmat.convertTo(rmat, CV_64FC1);
    tvec.convertTo(tvec, CV_64FC1);
    Eigen::Matrix3d rotated_matrix;  // 旋转矩阵
    Eigen::Vector3d transfer_vector; // 平移向量
    // 类型转换
    cv2eigen(rmat, rotated_matrix);
    cv2eigen(tvec, transfer_vector);
    // 获取欧拉角
    Eigen::Vector3d euler_angles = rotated_matrix.eulerAngles(0, 1, 2);

    // PNP 角度信息
    pnp_message.pitch = Rad2Deg(euler_angles[0]);
    pnp_message.yaw = Rad2Deg(euler_angles[1]);
    pnp_message.roll = Rad2Deg(euler_angles[2]);
    // PNP 距离信息
    pnp_message.distance = sqrt(transfer_vector.transpose() * transfer_vector);
    tvec.convertTo(tvec, CV_32FC1);
    pnp_message.tran_vec = tvec;
    pnp_message.rotat_vec = rvec;
    return pnp_message;
}

/**
 * @brief 计算神符真实转角
 * @param tracker 待预测追踪器对象
 * @return 实际角度
 */
float RunePredict::calculateRealAngle(tracker_ptr &tracker) {
    if (tracker->getFeatures().size() >= 2) {
        if (tracker->getFeatures().front()->getAngle() > 135.f &&
            tracker->getFeatures()[1]->getAngle() < -135.f) {
            this->round--;
        } else if (tracker->getFeatures().front()->getAngle() < -135.f &&
                   tracker->getFeatures()[1]->getAngle() > 135.f) {
            this->round++;
        }
        return tracker->getFeatures().front()->getAngle() + 360.f * this->round;
    } else {
        return tracker->getFeatures().front()->getAngle();
    }
}

/**
 * @brief 更新神符信息
 * @param tracker 待预测追踪器对象
 */
void RunePredict::updateRuneInfo(tracker_ptr &tracker) {
    // 更新角度
    this->real_angle.emplace_front(this->calculateRealAngle(tracker));
    // 更新半径
    this->radius = getDistance(tracker->getGoalFeature()->getSon1_Features()->getCenter(), getGoalCenter(tracker));
    if (this->real_angle.size() > (size_t) R_param.RECORD_FRAMES) {
        this->real_angle.pop_back();
    }
}

/**
 * @brief 更新预测信息
 * @param predict_matrix  预测信息矩阵
 * @param tracker 待预测追踪器对象
 */
void RunePredict::updatePredict(const cv::Matx21f &predict_matrix, tracker_ptr &tracker) {
    this->image_predict_angle = predict_matrix(0);
    float revise_angle = Deg2Rad(this->image_predict_angle);
    // 图像中的预测点
    tracker->setPredictedCenter().x = getGoalCenter(tracker).x + this->radius * cosf(revise_angle);
    tracker->setPredictedCenter().y = getGoalCenter(tracker).y - this->radius * sinf(revise_angle);
    // 云台预测角度
    tracker->setPredictedAngle() = A_param.predictConfidence * calculateRelativeAngle(cameraParam.cameraMatrix,
                                                                                      cameraParam.distCoeff,
                                                                                      tracker->setPredictedCenter()) +
                                   (1 - A_param.predictConfidence) * tracker->getGoalFeature()->getRelativeAngles();
}

/**
 * @brief 更新卡尔曼信息
 * @param tracker 待预测追踪器对象
 */
void RunePredict::updateKalman(tracker_ptr &tracker, float flying_time) {
    if (tracker->getKalmanInfo() && tracker->getFeatures().size() >= 3 && time_sequence.size() >= 3) {
        float total_time = static_cast<float >(time_sequence[0] - time_sequence[2]) / (2.3 * CLOCKS_PER_SEC);
        // 设置状态转移矩阵
        float predict_time = flying_time * 15;
        tracker->setRuneKF().setA(Matx<float, 2, 2>{1, predict_time, 0, 1});
        // 判断是否闪烁
        if (!tracker->isShine()) {
            // 更新角速度
            speed = (this->real_angle.front() -
                     this->real_angle[2]) /
                    total_time;
            // 更新
            tracker->setRuneKF().correct({this->real_angle.front(), speed});
        }
        // 预测
        auto predictVec = tracker->setRuneKF().predict();
        this->updatePredict(predictVec, tracker);
    } else {
        cv::Matx21f initVec = cv::Matx21f(this->real_angle.front(), 0);
        tracker->setRuneKF().init(initVec, 1e-2);
        this->updatePredict(initVec, tracker);
        tracker->setKalmanInfo() = true;
    }
}

/**
 * @brief 对外提供预测接口
 * @param tracker 待预测追踪器对象
 */
void RunePredict::predict(tracker_ptr &tracker, Compensate_ptr &compensate) {
    this->updateRuneInfo(tracker);
    this->updatePnP(tracker);
    compensate->compensate(tracker, _gyroData.front());
    this->updateKalman(tracker, compensate->getTime());
}

